import tensorflow
import numpy as np

'''
最大池化层
'''


class MaxPool(object):

    def __init__(self, x, kernel, strides, padding="VALID", active_fn=None, tf=tensorflow, d_type=np.float32):
        self.tf = tf
        self.x = x
        self.padding = padding
        self.kernel = kernel
        self.strides = strides
        self.activation_fun = active_fn

        self.x_max_pool = self.tf.nn.max_pool(self.x, self.kernel, self.strides, self.padding)

        if active_fn is None:
            self.y = self.x_max_pool
        else:
            self.y = active_fn(self.x_max_pool)


'''
均值池化
'''


class AvgPool(object):

    def __init__(self, x, kernel, strides, padding="VALID", active_fn=None, tf=tensorflow, d_type=np.float32):
        self.tf = tf
        self.x = x
        self.padding = padding
        self.kernel = kernel
        self.strides = strides
        self.activation_fun = active_fn

        self.x_avg_pool = self.tf.nn.avg_pool(self.x, self.kernel, self.strides, self.padding)

        if active_fn is None:
            self.y = self.x_avg_pool
        else:
            self.y = active_fn(self.x_avg_pool)


'''
二维卷积
'''


class Convolution2D(object):

    def __call__(self, x, f, s, p):
        return self.tf.nn.conv2d(x, f, s, p)

    def __init__(self, x, kernel, strides, padding="VALID", active_fn=None, tf=tensorflow, d_type=np.float32):
        self.tf = tf
        self.x = x
        self.padding = padding
        self.kernel = kernel
        self.strides = strides
        self.activation_fun = active_fn
        self.Weights = self.tf.Variable(self.tf.truncated_normal(kernel, stddev=0.1, dtype=d_type),
                                        dtype=d_type)
        self.Biases = self.tf.Variable((self.tf.constant(0.1, d_type, [kernel[3]])), dtype=d_type)

        self.x_convolution_Weights = self.tf.nn.conv2d(self.x, self.Weights, self.strides, self.padding)
        self.x_convolution_Weights_plus_Biases = self.x_convolution_Weights + self.Biases
        if active_fn is None:
            self.y = self.x_convolution_Weights_plus_Biases
        else:
            self.y = active_fn(self.x_convolution_Weights_plus_Biases)
